# Prepare dataset
import h5py, torch
from torch.utils.data import DataLoader, TensorDataset  
from sklearn.model_selection import train_test_split

import HighLevelFeatures as HLF
import numpy as np
import io
from PIL import Image
from evaluate import plot_Etot_Einc 

import pytorch_lightning as pl  
import torch.nn.functional as F
from pytorch_lightning.loggers import TensorBoardLogger  
from pytorch_lightning.callbacks import EarlyStopping

class ARG:
    def __init__(self):
        self.mode = 'all'
        self.buf = io.BytesIO()

def evaluate(x, y, y_hat, tensorboard, current_epoch):
    # avg draw
    y_hat = y_hat.detach().cpu().numpy()
    x = x.detach().cpu().numpy()
    y = y.detach().cpu().numpy()

    hlf = HLF.HighLevelFeatures('photon', filename='binning_dataset_1_photons.xml')
    hlf_ref = HLF.HighLevelFeatures('photon', filename='binning_dataset_1_photons.xml')

    _buf = io.BytesIO()
    hlf.DrawAverageShower(y_hat, _buf, 'Photon Average Shower')
    _buf.seek(0)
    tensorboard.add_image('Photon Average Shower', np.array(Image.open(_buf)), current_epoch, dataformats='HWC')

    # hist
    hlf.CalculateFeatures(y)
    hlf.Einc = x
    hlf_ref.CalculateFeatures(y_hat)
    hlf_ref.Einc = x

    arg = ARG()
    chi2 = plot_Etot_Einc(hlf, hlf_ref, arg)
    arg.buf.seek(0)
    tensorboard.add_scalar('Etot_Einc_Chi2', chi2, current_epoch)
    tensorboard.add_image('Etot_Einc', np.array(Image.open(arg.buf)), current_epoch, dataformats='HWC')


class SimpleMLP(pl.LightningModule):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = torch.nn.Linear(50, 100)
        self.fc2 = torch.nn.Linear(100, 200)
        self.fc4 = torch.nn.Linear(200, 368)
        self.relu = torch.nn.ReLU()
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc4(x)
        x = self.relu(x)
        return x
    
    def configure_optimizers(self):  
        return torch.optim.SGD(self.parameters(), lr=0.01, momentum=0.9)  
    
    def training_step(self, batch, batch_idx):  
        x, y = batch  
        y_hat = self(x)  
        loss = F.cross_entropy(y_hat, y)
        self.log('train_loss', loss, on_step=True, on_epoch=True)
        return loss  
    
    def validation_step(self, batch, batch_idx):  
        x, y = batch  
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('val_loss', loss, on_step=True, on_epoch=True)
        if batch_idx == 0:
            evaluate(x, y, y_hat, self.logger.experiment, self.current_epoch)

    def predict_step(self, batch, batch_idx, dataloader_idx=None):  
        x = batch[0]
        y_hat = self(x)
        return y_hat


photon_file = h5py.File('../data/dataset_1_photons_1.hdf5', 'r')

energies = photon_file['incident_energies'] # shape [121000, 1]
showers = photon_file['showers'] # shape [121000, 368]

print(energies.shape, showers.shape)

# Split the dataset into training and testing

X_train, X_test, y_train, y_test = train_test_split(energies,showers, test_size=0.2, random_state=114514, shuffle=False)

print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)

# make it Data loader
train_dataset = TensorDataset(torch.tensor(X_train, dtype=torch.double), torch.tensor(y_train, dtype=torch.double))
test_dataset = TensorDataset(torch.tensor(X_test, dtype=torch.double), torch.tensor(y_test, dtype=torch.double))

train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True, num_workers=3)
test_loader = DataLoader(test_dataset, batch_size=1024, shuffle=True, num_workers=3)


# torch.set_float32_matmul_precision('high')
tensorboard_logger = TensorBoardLogger(save_dir='.')  
model = SimpleMLP()
early_stopping = EarlyStopping('val_loss_epoch', patience=10)
trainer = pl.Trainer(max_epochs=1000, devices='auto', accelerator='auto', logger=tensorboard_logger, callbacks=[early_stopping])

trainer.fit(model, train_loader, test_loader)


# torch.save(model.state_dict(), 'model/model-MLP.pth')
# # load and test the model
# from torch.utils.data import DataLoader, TensorDataset  
# import torch

# model_test = SimpleMLP()
# model_test.load_state_dict(torch.load('model/model-MLP.pth'))

# energies_dataset = TensorDataset(torch.tensor(energies, dtype=torch.float))
# energies_dataset_loader = DataLoader(energies_dataset, batch_size=64, shuffle=False, num_workers=3)
# showers_predict = trainer.predict(model_test, dataloaders=energies_dataset_loader)
# reshaped_showers_predict = torch.cat(showers_predict, dim=0)

# with h5py.File("../data/eval-1-photon.hdf5", "w") as f:
#     f.create_dataset('incident_energies', data=energies)
#     f.create_dataset('showers', data=reshaped_showers_predict)